import os
import torchvision
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10, MNIST, FashionMNIST,VisionDataset
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
import cv2
import random


class drink(VisionDataset):
    def __init__(self,root,classes=range(3),train=True,
                            percent = [1]*3,
                            transform=None):
        kaiwei = os.path.join(root,"kaiwei")
        nongfu = os.path.join(root,"nongfu")
        moli = os.path.join(root,"moli")
        self.train = train
        self.transform = transform
        self.percent = percent
        
        if self.train:
            self.train_data = []
            self.train_labels = []
        else:
            self.test_data = []
            self.test_labels = []
        
        if 0 in classes:
            self.load_data(kaiwei,0)
        if 1 in classes:
            self.load_data(nongfu,1)
        if 2 in classes:
            self.load_data(moli,2)

    def load_data(self,path,label):
        images = os.listdir(path)
        indexes = list(range(len(images)))
        random.shuffle(indexes)
        if self.train:
            for i in range(len(indexes)):
                if i<0.9*len(indexes):
                    if(random.random()<self.percent[label]):
                        img = cv2.imread(os.path.join(path,images[i]))
                        img = cv2.resize(img,(128,128))
                        self.train_data.append(img)
                        self.train_labels.append(label)
                else:
                    break
        else:
            for i in range(len(indexes)):
                img = cv2.imread(os.path.join(path,images[i]))
                img = cv2.resize(img,(128,128))
                if i<0.1*len(indexes):
                    self.test_data.append(img)
                    self.test_labels.append(label)
                else:
                    break
    
    def __len__(self):
        if self.train:
            return len(self.train_data)
        else:
            return len(self.test_data)
    
    def __getitem__(self,index):
        if self.train:
            img, target = self.train_data[index], self.train_labels[index]
            img = Image.fromarray(img)
            if self.transform is not None:
                img = self.transform(img)

            return img, target
        else:
            img, target = self.test_data[index], self.test_labels[index]
            img = Image.fromarray(img)

            if self.transform is not None:
                img = self.transform(img)
            return img, target

    def sum(self):
        data_num = [0]*3
        if(self.train):
            for i in range(len(self.train_labels)):
                data_num[self.train_labels[i]]+=1
        else:
            for i in range(len(self.test_labels)):
                data_num[self.test_labels[i]]+=1
        print(data_num)

        
    

class cifar10(CIFAR10):
    def __init__(self, root,
                 classes=range(10),
                 percent = [1]*10,
                 train=True,
                 transform=None,
                 target_transform=None,
                 download=False):
        super(cifar10, self).__init__(root,
                                       train=train,
                                       transform=transform,
                                       target_transform=target_transform,
                                       download=download)
        
        self.train = train
        self.num_classes = len(classes)

        # Select subset of classes
        if self.train:
            self.train_data = []
            self.train_labels = []

            for i in range(len(self.data)):
                if self.targets[i] in classes:
                    if(random.random()<percent[self.targets[i]]):
                        self.train_data.append(self.data[i])
                        self.train_labels.append(int(self.targets[i]))
                    
        
        else:
            self.test_data = []
            self.test_labels = []

            for i in range(len(self.data)):
                if self.targets[i] in classes:
                    self.test_data.append(self.data[i])
                    self.test_labels.append(int(self.targets[i]))



    def __getitem__(self, index):
        if self.train:
            img, target = self.train_data[index], self.train_labels[index]
            img = Image.fromarray(img)

            if self.transform is not None:
                img = self.transform(img)

            if self.target_transform is not None:
                target = self.target_transform(target)
            return img, target
        else:
            img, target = self.test_data[index], self.test_labels[index]
            img = Image.fromarray(img)

            if self.transform is not None:
                img = self.transform(img)

            if self.target_transform is not None:
                target = self.target_transform(target)
            return img, target

    def __len__(self):
        if self.train:
            return len(self.train_data)
        else:
            return len(self.test_data)

    def sum(self):
        data_num = [0]*10
        if(self.train):
            for i in range(len(self.train_labels)):
                data_num[self.train_labels[i]]+=1
        else:
            for i in range(len(self.test_labels)):
                data_num[self.test_labels[i]]+=1
        print(data_num)

def Dataset(dataset_name,batch_size=128, classes=range(10),percent = [1]*10, data_dir='data'):

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    if dataset_name == 'cifar10':
        print('Dataset: CIFAR10.')
        trainset = cifar10(root=os.path.join(data_dir, 'cifar10'), train=True, download=True, transform=transforms.Compose([
                    transforms.RandomHorizontalFlip(),
                    transforms.RandomCrop(32, 4),
                    transforms.ToTensor(),
                    normalize]),
                    classes=classes,
                    percent = percent)

        testset = cifar10(root=os.path.join(data_dir, 'cifar10'), train=False, download=True, transform=transforms.Compose([
                    transforms.ToTensor(),
                    normalize]),
                    classes=classes)
    elif dataset_name == 'drink':
        trainset = drink(root=os.path.join(data_dir, 'drink'), train=True, transform=transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(100),
            transforms.ToTensor()]),
            percent = percent,
            classes=classes)

        testset = drink(root=os.path.join(data_dir, 'drink'), train=False, transform=transforms.Compose([
                    transforms.ToTensor()]),
                    classes=classes)

    else:    
        raise NotImplementedError("Only cifar10, imagenet, fashion mnist are allowed.")
    print("trainset num:")
    trainset.sum()
    print("testset num:")
    testset.sum()
    trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=12)
    testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=12)
    return trainloader, testloader
